/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: op execute.
 * Author: qingjicheng
 * Create: 2023-3-24
 */

#include "op_executor.h"
#include "transport_heterog_roce_pub.h"
#include "transport_roce_pub.h"
#include "transport_shm_event_pub.h"

namespace hccl {
HcclResult OpExecutor::Send(TransTask *taskPtr, TransportHeterog *transportPtr) const
{
    HcclRequestInfo* request = nullptr;
    void *sendMem = taskPtr->sendMem;
    TransData sendData(reinterpret_cast<u64>(sendMem), reinterpret_cast<u64>(nullptr), taskPtr->count,
        taskPtr->dataType, taskPtr->errorFlag);
    TransportEndPointInfo srcEp(0, taskPtr->srcRank, taskPtr->tag);
    TransportEndPointInfo dstEp(0, taskPtr->dstRank, taskPtr->tag);
    TransportEndPointParam epParam(srcEp, dstEp);
    CHK_RET(transportPtr->Isend(sendData, epParam, request));
    s32 flag = HCCL_TEST_INCOMPLETED;
    HcclStatus compState;
    while (flag != HCCL_TEST_COMPLETED) {
        CHK_RET(transportPtr->Test(*request, flag, compState));
        if (flag != HCCL_TEST_COMPLETED) {
            SaluSleep(LOOP_SLEEP_TIME_US);
        }
    }
    return HCCL_SUCCESS;
}

HcclResult OpExecutor::Isend(TransTask *taskPtr, TransportHeterog *transportPtr, HcclRequestInfo *&request) const
{
    TransData sendData(reinterpret_cast<u64>(taskPtr->sendMem), reinterpret_cast<u64>(nullptr), taskPtr->count,
        taskPtr->dataType);
    TransportEndPointInfo srcEp(0, taskPtr->srcRank, taskPtr->tag);
    TransportEndPointInfo dstEp(0, taskPtr->dstRank, taskPtr->tag);
    TransportEndPointParam epParam(srcEp, dstEp);

    TransportHeterogRoce *transportHeterogRoce = dynamic_cast<TransportHeterogRoce *>(transportPtr);
    // ES拉远推理是TransportRoce
    TransportRoce *transportRoce = dynamic_cast<TransportRoce *>(transportPtr);

    // ES训练的PS响应流程
    if ((transportHeterogRoce != nullptr) && (transportRoce == nullptr)) {
        HCCL_DEBUG("transportPtr is TransportHeterogRoce in ps Isend, responseBuffer.addr[%#llx].",
            taskPtr->responseBuffer.addr);

        sendData.dstBuf = taskPtr->responseBuffer.addr;
        HcclEnvelope envelope(0, sendData, epParam, taskPtr->responseBuffer.buffKey, 0);

        CHK_RET(transportPtr->Iwrite(sendData, envelope, request));
    } else {
        CHK_RET(transportPtr->Isend(sendData, epParam, request));
    }

    return HCCL_SUCCESS;
}

void OpExecutor::GetProbeMsgInfo(TransportHeterog *heterogPtr, HcclMessageInfo *msg,
    u32 &recvDataType, HcclUserRequire &userRequire)
{
    TransportShmEvent *transportShm = dynamic_cast<TransportShmEvent *>(heterogPtr);
    if (transportShm != nullptr) {
        recvDataType = msg->envelope.pcieEnvelope.dataType;
        userRequire.tableId = msg->envelope.pcieEnvelope.tableId;
        userRequire.globalStep = msg->envelope.pcieEnvelope.globalStep;
    } else {
        recvDataType = msg->envelope.envelope.transData.dataType;
        userRequire.tableId = msg->envelope.envelope.transData.tableId;
        userRequire.globalStep = msg->envelope.envelope.transData.globalStep;
    }
}

void OpExecutor::SetTaskOutputRegStatusInfo(TransTask *taskPtr,  HcclStatus &status, u32 &recvDataType,
    HcclUserRequire &userRequire)
{
    if (!taskPtr->outputs.empty()) {
        std::vector<void *> &outputs = taskPtr->outputs;

        u32 dataTypeBytes = SIZE_TABLE[recvDataType];
        ReqStatus *reqStatus = static_cast<ReqStatus *>(outputs[OUTPUT_REQ_STATUS_IDX]);

        if (dataTypeBytes % sizeof(u64) != 0) {
            // ES推理会走该分支
            reqStatus->actualSize += (status.count / sizeof(u64));
        } else {
            u32 factor = dataTypeBytes / sizeof(u64);
            // 当前在特征准入时，返回的key数量包含counter，为实际key数量的两倍
            reqStatus->actualSize += (status.count * factor);
        }

        reqStatus->tableId = userRequire.tableId;
        reqStatus->globalStep = userRequire.globalStep;
        HCCL_INFO("TaskImprobeRecv success, key actual size[%d], DataType[%s], bytesNum[%u], tableId[%u], tag[%d]",
            reqStatus->actualSize, GetDataTypeEnumStr(recvDataType).c_str(), dataTypeBytes, reqStatus->tableId,
            reqStatus->tag);
        PRINT_ARRAY(static_cast<s64 *>(taskPtr->recvMem), reqStatus->actualSize, "PSRecvedKeys");
    }
}
} // namespace hccl